from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd
from collections import deque

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
stock_hfq_df = ak.stock_zh_a_hist(symbol="603777", adjust="hfq").iloc[:, :6]
# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])

class MACDBacktest(bt.Strategy):
    params = (
        ("macd1_period", 12),
        ("macd2_period", 26),
        ("signal_period", 9),
        ("divergence_threshold", 0.01)
    )

    def __init__(self):
        self.macd = bt.indicators.MACD(
            self.data,
            period_me1=self.p.macd1_period,
            period_me2=self.p.macd2_period,
            period_signal=self.p.signal_period
        )

    def next(self):
        if self.macd.macd[0] > self.macd.signal[0] and \
           self.macd.macd[-1] < self.macd.signal[-1] and \
           self.macd.macd[0] < self.macd.macd[-1] and \
           (self.macd.macd[-1] - self.macd.macd[0]) / self.macd.macd[-1] > self.p.divergence_threshold:
            self.buy()
        elif self.macd.macd[0] < self.macd.signal[0] and \
             self.macd.macd[-1] > self.macd.signal[-1] and \
             self.macd.macd[0] > self.macd.macd[-1] and \
             (self.macd.macd[0] - self.macd.macd[-1]) / self.macd.macd[-1] > self.p.divergence_threshold:
            self.sell()


# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(1991, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MACDBacktest)

# 设置初始资金和手续费
start_cash = 1000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")

# 绘制图表
cerebro.plot(style='candlestick')